﻿using System;
using System.IO;
using System.Runtime.InteropServices;
using System.Threading;
using Ast.Core.Events;
using Ast.Core.Pe;
using Ast.Core.Pe.Definitions;
using Ast.Core.Services.Extensions;
using Ast.Framework;
using Castle.Core.Logging;

namespace Ast.Core.Services
{
    public interface IKernelReloadService : ILockable, ISingletonDependency
    {
        IntPtr NewKernel { get; }
        void Reload();
    }

    public class KernelReloadService : IKernelReloadService, IDisposable
    {
        private readonly IPeService _peService;
        private readonly IDeviceController _deviceController;
        private readonly IPsLoadedModuleListService _psLoadedModuleListService;
        private readonly ISsdtService _ssdtService;
        private readonly IKernelReloadEvents _kernelReloadEvents;

        private readonly object _syncRoot = new object();
        public IntPtr NewKernel { get; private set; }

        public KernelReloadService(IPeService peService, IDeviceController deviceController, IPsLoadedModuleListService psLoadedModuleListService, ISsdtService ssdtService, IKernelReloadEvents kernelReloadEvents)
        {
            _peService = peService;
            _deviceController = deviceController;
            _psLoadedModuleListService = psLoadedModuleListService;
            _ssdtService = ssdtService;
            _kernelReloadEvents = kernelReloadEvents;

            NewKernel = IntPtr.Zero;

            Logger = NullLogger.Instance;
        }

        public ILogger Logger { get; set; }

        public void Reload()
        {
            RelaseKernel();

            // events
            _kernelReloadEvents.OnRealoding();

            // get kernel file path
            var ntModule = _psLoadedModuleListService.GetNtModule();
            if (ntModule == null)
            {
                Logger.Error("failed to load nt module");
                return;
            }
            var currentKernelBase = new IntPtr((int)ntModule.StartAddress);
            var kernelFilePath = ntModule.Path;
            kernelFilePath = Path.Combine(Environment.SystemDirectory, Path.GetFileName(kernelFilePath));
            Logger.InfoFormat("current kernel image base 0x{0:x}, file {1}", currentKernelBase.ToInt32(), kernelFilePath);

            // load pe
            Logger.InfoFormat("loading kernel file at {0}", kernelFilePath);
            var peModule = _peService.Load(kernelFilePath);
            if (peModule == null)
            {
                Logger.Error("failed to load pe image");
                return;
            }

            // allocate memory
            var memoryBlock = _deviceController.AllocateMemory((int) peModule.OptionalHeader.SizeOfImage);
            var address = memoryBlock.Value;
            if (address == IntPtr.Zero)
            {
                Logger.Error("failed to allocate memory for image");
                // return;
            }

            // DosHeader
            _deviceController.WriteAs(address, peModule.DosHeader);

            // NewExeFileHeader
            address += (int) peModule.DosHeader.NewExeHeaderOffset;
            _deviceController.WriteAs(address, peModule.NewExeFileHeader);

            // OptionalHeader
            address += Marshal.SizeOf<NewExeFileHeader>();
            if (peModule.ModuleType == PeModuleType.Pe32)
            {
                _deviceController.WriteAs(address, (OptionalHeader32) peModule.OptionalHeader);
            }
            else
            {
                _deviceController.WriteAs(address, (OptionalHeader64) peModule.OptionalHeader);
            }

            // Sections
            address += peModule.NewExeFileHeader.SizeOfOptionalHeader;
            for (var index = 0; index < peModule.SectionHeaders.Length; index++)
            {
                var sectionHeader = peModule.SectionHeaders[index];
                _deviceController.WriteAs(address + index*SectionHeader.StructSize, sectionHeader);
            }

            // Sections' Data
            foreach (var sectionContent in peModule.SectionContents)
            {
                address = memoryBlock.Value + (int) sectionContent.Key.VirtualAddress;
                Logger.InfoFormat("writing section {0} to 0x{1:x}, size 0x{2:x}, virtual address 0x{3:x}", sectionContent.Key.Name, address.ToInt32(), sectionContent.Value.Length, sectionContent.Key.VirtualAddress);
                _deviceController.WriteBytes(address, sectionContent.Value);
            }

            // fix relocation blocks
            FixRelocations(peModule, memoryBlock.Value, currentKernelBase);

            // fix ssdt
            FixSddt(memoryBlock.Value, currentKernelBase);

            NewKernel = memoryBlock.Value;
            Logger.InfoFormat("reloaded kernel at 0x{0:x}, original at 0x{1:x}, offset 0x{2:x}", memoryBlock.Value.ToInt32(), currentKernelBase.ToInt32(), memoryBlock.Value.ToInt32() - currentKernelBase.ToInt32());

            // events
            _kernelReloadEvents.OnReloaded(currentKernelBase, NewKernel);
        }

        private void FixSddt(IntPtr newImageBase, IntPtr currentImageBase)
        {
            if (!_ssdtService.FindAddress())
            {
                Logger.Error("failed to search ssdt");
                return;
            }

            Logger.Info("fixing ssdt");
            var offset = newImageBase.ToInt32() - currentImageBase.ToInt32();
            var newSsdt = _ssdtService.Ssdt + offset;
            var curTbl = _deviceController.ReadAs<SsdtService.KServiceTableDescriptor>(_ssdtService.Ssdt);
            var newTbl = curTbl; // copy value
            newTbl.Ntoskrnl.ServiceTableBase += offset;
            newTbl.Ntoskrnl.ParamTableBase += offset;
            _deviceController.WriteAs(newSsdt, newTbl); // fixing table

            for (var i = 0; i < newTbl.Ntoskrnl.NumberOfService; i++)
            {
                var adr = newTbl.Ntoskrnl.ServiceTableBase + i*4;
                var curFunc = _deviceController.ReadAsIntPtr(adr).Value;
                var newFunc = curFunc + offset;
                _deviceController.WriteAs(adr, newFunc); // fixing function
            }

            Logger.InfoFormat("fixed {0} ssdt functions, new ssdt 0x{1:x}, old ssdt 0x{2:x}", newTbl.Ntoskrnl.NumberOfService, newSsdt.ToInt32(), _ssdtService.Ssdt.ToInt32());
        }

        private void FixRelocations(PeModule peModule, IntPtr newImageBase, IntPtr curImageBase)
        {
            if (!peModule.RelocationDirectoryEntry.HasValue) return;

            var sectionHeader = peModule.FindSectionContainingRva(peModule.RelocationDirectoryEntry.Value.VirtualAddress);
            if (!sectionHeader.HasValue) return;

            Logger.Info("fixing relocations");
            using (var memoryStream = new MemoryStream(peModule.SectionContents[sectionHeader.Value]))
            {
                var expectedImageBase = (uint) peModule.OptionalHeader.ImageBase;
                var offset = (uint) curImageBase.ToInt32() - expectedImageBase;
                Logger.InfoFormat("current image base 0x{0:x}, expected base 0x{1:x}, offset 0x{2:x}", curImageBase.ToInt32(), expectedImageBase, offset);

                var count = 0;
                var relocationDirectoryDescriptor = memoryStream.ReadAs<RelocationDirectoryDescriptor>();
                while (relocationDirectoryDescriptor.VirtualAddress != 0)
                {
                    var relocationTypeOffsets = memoryStream.ReadArrayAs<RelocationTypeOffset>(relocationDirectoryDescriptor.Count);
                    foreach (var relocationTypeOffset in relocationTypeOffsets)
                    {
                        if (relocationTypeOffset.Type != RelocationType.ImageRelBasedHighlow) continue;

                        var addressToRelocate = new IntPtr((int) ((uint) newImageBase.ToInt32() + relocationDirectoryDescriptor.VirtualAddress + relocationTypeOffset.Offset));
                        var currentVal = (uint) _deviceController.ReadAsIntPtr(new IntPtr((int) addressToRelocate)).Value.ToInt32();
                        var newVal = currentVal + offset;
                        // Logger.InfoFormat("fixing address 0x{0:x}, current value 0x{1:x}, new value 0x{2:x}", addressToRelocate.ToInt32(), currentVal, newVal);
                        _deviceController.WriteAs(addressToRelocate, new IntPtr((int) newVal));
                        count++;
                    }
                    relocationDirectoryDescriptor = memoryStream.ReadAs<RelocationDirectoryDescriptor>(); // next
                }
                Logger.InfoFormat("fixed {0} relocations", count);
            }
        }

        public void AcquireLock()
        {
            Monitor.Enter(_syncRoot);
        }

        public void ReleaseLock()
        {
            Monitor.Exit(_syncRoot);
        }

        private void RelaseKernel(bool isDisposing = false)
        {
            AcquireLock();
            if (NewKernel != IntPtr.Zero)
            {
                // events
                if (!isDisposing) _kernelReloadEvents.OnUnloading(NewKernel);

                _deviceController.FreeMemory(NewKernel);
                NewKernel = IntPtr.Zero;
                Logger.Info("released kernel");

                // events
                if (!isDisposing) _kernelReloadEvents.OnUnloaded();
            }
            ReleaseLock();
        }

        public void Dispose()
        {
            RelaseKernel(true);
        }
    }
}